use crate::ast::*;
use crate::error::{CompilerError, Result};
use inkwell::values::BasicValueEnum;

use super::LLVMBackend;

impl<'ctx> LLVMBackend<'ctx> {
    /// Generate for statement
    pub(super) fn generate_for_statement(&mut self, for_stmt: &ForStmt) -> Result<()> {
        match &for_stmt.variant {
            crate::ast::stmt::ForVariant::Traditional {
                initializer,
                condition,
                increment,
            } => {
                // Generate initializer if present
                if let Some(init) = initializer {
                    self.generate_statement(init)?;
                }

                // Create loop blocks
                let cond_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "for.cond");
                let body_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "for.body");
                let inc_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "for.inc");
                let end_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "for.end");

                // Jump to condition block
                let _ = self.builder.build_unconditional_branch(cond_block);

                // Generate condition block
                self.builder.position_at_end(cond_block);
                if let Some(cond) = condition {
                    let condition = self.generate_expression(cond)?;
                    let condition_int = match condition {
                        BasicValueEnum::IntValue(int_val) => int_val,
                        _ => {
                            return Err(CompilerError::internal("For condition must be an integer"))
                        }
                    };

                    // Convert integer to boolean (i1) if necessary
                    let condition_bool = if condition_int.get_type() == self.context.bool_type() {
                        // Already i1, use directly
                        condition_int
                    } else {
                        // Convert i32/other int to i1
                        let zero = condition_int.get_type().const_int(0, false);
                        self.builder.build_int_compare(
                            inkwell::IntPredicate::NE,
                            condition_int,
                            zero,
                            "cond_bool",
                        )?
                    };

                    let _ = self.builder.build_conditional_branch(
                        condition_bool,
                        body_block,
                        end_block,
                    );
                } else {
                    let _ = self.builder.build_unconditional_branch(body_block);
                }

                // Generate body block
                self.builder.position_at_end(body_block);
                self.enter_scope();
                self.generate_statement(&for_stmt.body)?;
                self.exit_scope()?;
                let _ = self.builder.build_unconditional_branch(inc_block);

                // Generate increment block
                self.builder.position_at_end(inc_block);
                if let Some(inc) = increment {
                    self.generate_statement(inc)?;
                }
                let _ = self.builder.build_unconditional_branch(cond_block);

                // Position at end block
                self.builder.position_at_end(end_block);
            }
            crate::ast::stmt::ForVariant::Condition { condition } => {
                // Create loop blocks
                let cond_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "for.cond");
                let body_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "for.body");
                let end_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "for.end");

                // Jump to condition block
                let _ = self.builder.build_unconditional_branch(cond_block);

                // Generate condition block
                self.builder.position_at_end(cond_block);
                let condition = self.generate_expression(condition)?;
                let condition_int = match condition {
                    BasicValueEnum::IntValue(int_val) => int_val,
                    _ => return Err(CompilerError::internal("For condition must be an integer")),
                };

                // Convert integer to boolean (i1) if necessary
                let condition_bool = if condition_int.get_type() == self.context.bool_type() {
                    // Already i1, use directly
                    condition_int
                } else {
                    // Convert i32/other int to i1
                    let zero = condition_int.get_type().const_int(0, false);
                    self.builder.build_int_compare(
                        inkwell::IntPredicate::NE,
                        condition_int,
                        zero,
                        "cond_bool",
                    )?
                };

                let _ =
                    self.builder
                        .build_conditional_branch(condition_bool, body_block, end_block);

                // Generate body block
                self.builder.position_at_end(body_block);
                self.enter_scope();
                self.generate_statement(&for_stmt.body)?;
                self.exit_scope()?;
                let _ = self.builder.build_unconditional_branch(cond_block);

                // Position at end block
                self.builder.position_at_end(end_block);
            }
            crate::ast::stmt::ForVariant::Range {
                variable,
                variable_type,
                iterable,
            } => {
                // Range-based for loop: for var v T in slice {}
                // Desugar to: for (var i = 0; i < slice.len(); i++) { var v = slice[i]; body }

                // 1. Evaluate the iterable expression (slice)
                let slice_value = self.generate_expression(iterable)?;

                // Get slice type to verify it's actually a slice
                let vector_type = self
                    .variable_type_map
                    .values()
                    .find(|t| matches!(t, Type::Vec(_)))
                    .cloned();

                if vector_type.is_none() {
                    return Err(CompilerError::internal("for-in requires a slice type"));
                }

                // Extract slice struct to get data pointer and length
                let slice_struct = slice_value.into_struct_value();
                let data_ptr = self
                    .builder
                    .build_extract_value(slice_struct, 0, "slice_data")?
                    .into_pointer_value();
                let slice_len = self
                    .builder
                    .build_extract_value(slice_struct, 1, "slice_len")?
                    .into_int_value();

                // 2. Create loop index variable (i64)
                let i64_type = self.context.i64_type();
                let index_ptr = self.builder.build_alloca(i64_type, "__for_index")?;
                let zero = i64_type.const_int(0, false);
                self.builder.build_store(index_ptr, zero)?;

                // 3. Create loop blocks
                let cond_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "for.cond");
                let body_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "for.body");
                let inc_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "for.inc");
                let end_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "for.end");

                // Jump to condition block
                self.builder.build_unconditional_branch(cond_block)?;

                // 4. Generate condition block: i < slice.len()
                self.builder.position_at_end(cond_block);
                let current_index = self
                    .builder
                    .build_load(i64_type, index_ptr, "__for_index_val")?
                    .into_int_value();
                let condition = self.builder.build_int_compare(
                    inkwell::IntPredicate::SLT,
                    current_index,
                    slice_len,
                    "for_cond",
                )?;
                self.builder
                    .build_conditional_branch(condition, body_block, end_block)?;

                // 5. Generate body block
                self.builder.position_at_end(body_block);

                // Create loop variable: var v T = slice[i]
                let element_llvm_type = self.nrc_type_to_llvm_type(&Some(variable_type.clone()))?;
                let var_ptr = self.builder.build_alloca(element_llvm_type, variable)?;

                // Load element from slice: data_ptr[current_index]
                let element_ptr = unsafe {
                    self.builder.build_gep(
                        element_llvm_type,
                        data_ptr,
                        &[current_index],
                        "element_ptr",
                    )?
                };
                let element_value =
                    self.builder
                        .build_load(element_llvm_type, element_ptr, "element")?;
                self.builder.build_store(var_ptr, element_value)?;

                // Register loop variable in variable map
                self.variable_map.insert(variable.clone(), var_ptr.into());
                self.variable_type_map
                    .insert(variable.clone(), variable_type.clone());

                // Generate loop body
                self.enter_scope();
                self.generate_statement(&for_stmt.body)?;
                self.exit_scope()?;

                // Remove loop variable from scope
                self.variable_map.remove(variable);
                self.variable_type_map.remove(variable);

                // Jump to increment block
                self.builder.build_unconditional_branch(inc_block)?;

                // 6. Generate increment block: i++
                self.builder.position_at_end(inc_block);
                let current_index = self
                    .builder
                    .build_load(i64_type, index_ptr, "__for_index_val")?
                    .into_int_value();
                let one = i64_type.const_int(1, false);
                let next_index =
                    self.builder
                        .build_int_add(current_index, one, "__for_index_next")?;
                self.builder.build_store(index_ptr, next_index)?;
                self.builder.build_unconditional_branch(cond_block)?;

                // 7. Position at end block
                self.builder.position_at_end(end_block);
            }
        }
        Ok(())
    }

    /// Generate while statement
    pub(super) fn generate_while_statement(&mut self, while_stmt: &WhileStmt) -> Result<()> {
        // Create loop blocks
        let cond_block = self
            .context
            .append_basic_block(self.current_function.unwrap(), "while.cond");
        let body_block = self
            .context
            .append_basic_block(self.current_function.unwrap(), "while.body");
        let end_block = self
            .context
            .append_basic_block(self.current_function.unwrap(), "while.end");

        // Jump to condition block
        let _ = self.builder.build_unconditional_branch(cond_block);

        // Generate condition block
        self.builder.position_at_end(cond_block);
        let condition = self.generate_expression(&while_stmt.condition)?;
        let condition_int = match condition {
            BasicValueEnum::IntValue(int_val) => int_val,
            _ => {
                return Err(CompilerError::internal(
                    "While condition must be an integer",
                ))
            }
        };

        // Convert integer to boolean (i1) if necessary
        let condition_bool = if condition_int.get_type() == self.context.bool_type() {
            // Already i1, use directly
            condition_int
        } else {
            // Convert i32/other int to i1
            let zero = condition_int.get_type().const_int(0, false);
            self.builder.build_int_compare(
                inkwell::IntPredicate::NE,
                condition_int,
                zero,
                "cond_bool",
            )?
        };

        let _ = self
            .builder
            .build_conditional_branch(condition_bool, body_block, end_block);

        // Generate body block
        self.builder.position_at_end(body_block);
        self.generate_statement(&while_stmt.body)?;
        let _ = self.builder.build_unconditional_branch(cond_block);

        // Position at end block
        self.builder.position_at_end(end_block);

        Ok(())
    }

    /// Generate block expression (returns the value of the final expression)
    pub(super) fn generate_block_expression(
        &mut self,
        block_expr: &BlockExpr,
    ) -> Result<BasicValueEnum<'ctx>> {
        // Execute all statements
        for stmt in &block_expr.statements {
            self.generate_statement(stmt)?;
        }

        // Return the final expression's value, or a default value if none
        if let Some(final_expr) = &block_expr.final_expr {
            self.generate_expression(final_expr)
        } else {
            // No final expression, return void/unit (0 as a placeholder)
            Ok(self.context.i32_type().const_int(0, false).into())
        }
    }

    /// Generate if expression (returns a value)
    pub(super) fn generate_if_expression(
        &mut self,
        if_expr: &IfExpr,
    ) -> Result<BasicValueEnum<'ctx>> {
        // Generate condition
        let condition = self.generate_expression(&if_expr.condition)?;

        // Create basic blocks
        let then_block = self
            .context
            .append_basic_block(self.current_function.unwrap(), "if.then");
        let else_block = self
            .context
            .append_basic_block(self.current_function.unwrap(), "if.else");
        let merge_block = self
            .context
            .append_basic_block(self.current_function.unwrap(), "if.merge");

        // Create conditional branch
        let condition_int = match condition {
            BasicValueEnum::IntValue(int_val) => int_val,
            _ => return Err(CompilerError::internal("If condition must be an integer")),
        };

        // Convert integer to boolean (i1) if necessary
        let condition_bool = if condition_int.get_type() == self.context.bool_type() {
            // Already i1, use directly
            condition_int
        } else {
            // Convert i32/other int to i1
            let zero = condition_int.get_type().const_int(0, false);
            self.builder.build_int_compare(
                inkwell::IntPredicate::NE,
                condition_int,
                zero,
                "cond_bool",
            )?
        };

        let _ = self
            .builder
            .build_conditional_branch(condition_bool, then_block, else_block);

        // Generate then branch
        self.builder.position_at_end(then_block);
        let then_value = self.generate_expression(&if_expr.then_branch)?;
        let then_end_block = self.builder.get_insert_block().unwrap();
        let _ = self.builder.build_unconditional_branch(merge_block);

        // Generate else branch
        self.builder.position_at_end(else_block);
        let else_value = if let Some(else_branch) = &if_expr.else_branch {
            self.generate_expression(else_branch)?
        } else {
            // No else branch, return default value (0)
            self.context.i32_type().const_int(0, false).into()
        };
        let else_end_block = self.builder.get_insert_block().unwrap();
        let _ = self.builder.build_unconditional_branch(merge_block);

        // Generate merge block with phi node
        self.builder.position_at_end(merge_block);
        let phi = self.builder.build_phi(then_value.get_type(), "if.result")?;
        phi.add_incoming(&[(&then_value, then_end_block), (&else_value, else_end_block)]);

        Ok(phi.as_basic_value())
    }

    /// Generate if statement
    pub(super) fn generate_if_statement(&mut self, if_stmt: &IfStmt) -> Result<()> {
        // Generate condition
        let condition = self.generate_expression(&if_stmt.condition)?;

        // Create basic blocks
        let then_block = self
            .context
            .append_basic_block(self.current_function.unwrap(), "if.then");
        let else_block = if if_stmt.else_branch.is_some() {
            Some(
                self.context
                    .append_basic_block(self.current_function.unwrap(), "if.else"),
            )
        } else {
            None
        };
        let end_block = self
            .context
            .append_basic_block(self.current_function.unwrap(), "if.end");

        // Create conditional branch
        let condition_int = match condition {
            BasicValueEnum::IntValue(int_val) => int_val,
            _ => return Err(CompilerError::internal("If condition must be an integer")),
        };

        // Convert integer to boolean (i1) if necessary
        let condition_bool = if condition_int.get_type() == self.context.bool_type() {
            // Already i1, use directly
            condition_int
        } else {
            // Convert i32/other int to i1
            let zero = condition_int.get_type().const_int(0, false);
            self.builder.build_int_compare(
                inkwell::IntPredicate::NE,
                condition_int,
                zero,
                "cond_bool",
            )?
        };

        let _ = self.builder.build_conditional_branch(
            condition_bool,
            then_block,
            else_block.unwrap_or(end_block),
        );

        // Generate then block
        self.builder.position_at_end(then_block);
        self.enter_scope();
        self.generate_statement(&*if_stmt.then_branch)?;
        self.exit_scope()?;
        let _ = self.builder.build_unconditional_branch(end_block);

        // Generate else block if present
        if let Some(else_branch) = &if_stmt.else_branch {
            if let Some(else_block) = else_block {
                self.builder.position_at_end(else_block);
                self.enter_scope();
                self.generate_statement(&**else_branch)?;
                self.exit_scope()?;
                let _ = self.builder.build_unconditional_branch(end_block);
            }
        }

        // Position builder at end block
        self.builder.position_at_end(end_block);

        Ok(())
    }

    /// Generate match statement
    pub(super) fn generate_match_statement(&mut self, match_stmt: &MatchStmt) -> Result<()> {
        // Generate the expression to match
        let match_value = self.generate_expression(&match_stmt.expr)?;

        // Create a single end block for the entire match statement
        let match_end_block = self
            .context
            .append_basic_block(self.current_function.unwrap(), "match.end");

        // For now, implement a simple version that doesn't create new basic blocks
        if match_stmt.arms.len() == 1 {
            // Single arm - just execute it directly
            let arm = &match_stmt.arms[0];
            self.generate_statement(&arm.body)?;
            // Branch to end block
            let _ = self.builder.build_unconditional_branch(match_end_block);
            // Position builder at end block
            self.builder.position_at_end(match_end_block);
            return Ok(());
        }

        // For multiple arms, implement a simple if-else chain
        for (i, arm) in match_stmt.arms.iter().enumerate() {
            let condition =
                self.generate_pattern_condition(match_value, &arm.pattern, &arm.guard)?;

            // Convert condition to boolean (i1) if necessary
            let condition_bool = if condition.get_type() == self.context.bool_type() {
                // Already i1, use directly
                condition
            } else {
                // Convert i32/other int to i1
                let zero = condition.get_type().const_int(0, false);
                self.builder.build_int_compare(
                    inkwell::IntPredicate::NE,
                    condition,
                    zero,
                    "condition_bool",
                )?
            };

            // Create basic blocks for this arm
            let then_block = self
                .context
                .append_basic_block(self.current_function.unwrap(), &format!("match.then.{}", i));
            let else_block = if i < match_stmt.arms.len() - 1 {
                Some(self.context.append_basic_block(
                    self.current_function.unwrap(),
                    &format!("match.else.{}", i),
                ))
            } else {
                None
            };

            // Create conditional branch
            let _ = self.builder.build_conditional_branch(
                condition_bool,
                then_block,
                else_block.unwrap_or(match_end_block),
            );

            // Generate then block
            self.builder.position_at_end(then_block);
            self.generate_statement(&arm.body)?;
            let _ = self.builder.build_unconditional_branch(match_end_block);

            // Generate else block if present
            if let Some(else_block) = else_block {
                self.builder.position_at_end(else_block);
                // Continue to next arm
            } else {
                // This is the last arm, position at end block
                self.builder.position_at_end(match_end_block);
                return Ok(());
            }
        }

        // Position builder at end block
        self.builder.position_at_end(match_end_block);
        Ok(())
    }

    /// Generate match expression (returns a value)
    pub(super) fn generate_match_expression(
        &mut self,
        match_expr: &MatchExpr,
    ) -> Result<BasicValueEnum<'ctx>> {
        if match_expr.arms.len() == 1 {
            // Single arm - just execute it directly
            let arm = &match_expr.arms[0];

            // Generate pattern condition which will also bind variables
            let match_value = self.generate_expression(&match_expr.expr)?;
            let _condition =
                self.generate_pattern_condition(match_value, &arm.pattern, &arm.guard)?;

            return self.generate_expression(&arm.body);
        }

        // Evaluate match value once
        let match_value = self.generate_expression(&match_expr.expr)?;

        // Create end block - we'll collect results here with PHI node
        let end_block = self
            .context
            .append_basic_block(self.current_function.unwrap(), "match.end");

        // Track (result_value, source_block) pairs for PHI node
        let mut phi_incoming: Vec<(BasicValueEnum<'ctx>, inkwell::basic_block::BasicBlock<'ctx>)> =
            Vec::new();

        // Generate all arms in sequence
        let mut current_else_block = None;

        for (i, arm) in match_expr.arms.iter().enumerate() {
            // Position at the current check block
            if let Some(else_block) = current_else_block {
                self.builder.position_at_end(else_block);
            }

            // Save current variable map state to restore after pattern binding
            let saved_var_map = self.variable_map.clone();
            let saved_type_map = self.variable_type_map.clone();

            // Generate pattern condition (this may bind variables, but we'll discard them)
            let condition =
                self.generate_pattern_condition(match_value, &arm.pattern, &arm.guard)?;
            let condition_bool = if condition.get_type() == self.context.bool_type() {
                condition
            } else {
                let zero = condition.get_type().const_int(0, false);
                self.builder.build_int_compare(
                    inkwell::IntPredicate::NE,
                    condition,
                    zero,
                    "condition_bool",
                )?
            };

            // Restore variable maps - we don't want pattern bindings to leak into the check block
            self.variable_map = saved_var_map;
            self.variable_type_map = saved_type_map;

            // Create then block for this arm
            let then_block = self
                .context
                .append_basic_block(self.current_function.unwrap(), &format!("match.arm.{}", i));

            // Create else block for next arm (if not last arm)
            let next_else_block = if i < match_expr.arms.len() - 1 {
                Some(self.context.append_basic_block(
                    self.current_function.unwrap(),
                    &format!("match.check.{}", i + 1),
                ))
            } else {
                None
            };

            // Branch based on condition
            if i == match_expr.arms.len() - 1 {
                // Last arm - if it doesn't match, the match is non-exhaustive
                // For now, branch to a panic block
                let panic_block = self
                    .context
                    .append_basic_block(self.current_function.unwrap(), "match.panic");
                self.builder
                    .build_conditional_branch(condition_bool, then_block, panic_block)?;

                // Generate panic block
                self.builder.position_at_end(panic_block);

                // Call exit(1) to panic
                let exit_fn = self.module.get_function("exit").unwrap();
                self.builder.build_call(
                    exit_fn,
                    &[self.context.i32_type().const_int(1, false).into()],
                    "exit_call",
                )?;
                self.builder.build_unreachable()?;

                // Continue with then block (position will be set below)
            } else {
                self.builder.build_conditional_branch(
                    condition_bool,
                    then_block,
                    next_else_block.unwrap(),
                )?;
            }

            // Generate arm body in then block
            self.builder.position_at_end(then_block);

            // Now bind pattern variables properly in this block
            let saved_var_map = self.variable_map.clone();
            let saved_type_map = self.variable_type_map.clone();
            let _ = self.generate_pattern_condition(match_value, &arm.pattern, &arm.guard)?;

            // Generate arm body expression
            let arm_result = self.generate_expression(&arm.body)?;

            // Restore variable maps before leaving this arm
            self.variable_map = saved_var_map;
            self.variable_type_map = saved_type_map;

            // Record the result and source block for PHI
            let result_block = self.builder.get_insert_block().unwrap();
            phi_incoming.push((arm_result, result_block));

            // Branch to end
            self.builder.build_unconditional_branch(end_block)?;

            current_else_block = next_else_block;
        }

        // Position at end block and create PHI node
        self.builder.position_at_end(end_block);

        // Create PHI node to merge results from all arms
        let result_type = phi_incoming[0].0.get_type();
        let phi = self.builder.build_phi(result_type, "match.result")?;

        for (value, block) in phi_incoming {
            phi.add_incoming(&[(&value, block)]);
        }

        Ok(phi.as_basic_value())
    }
}
